package com.xuyuan.javaagent;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtField;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 通过ClassFileTransformer接口，可以在类加载之前，重写字节码(校验注解失效)
 */
public class CustomChangeSuperTransformer implements ClassFileTransformer {
    /**
     * 待修改的父类的属性列表
     */
    public static final Map<String, List<String>> subFields = new HashMap<>();
    /**
     * 待修改的父类全限定类名
     */
    public static final List<String> parentClassNames = new ArrayList<>();
    /**
     * 待监控的子类的包
     */
    private static final List<String> monitorPackages = new ArrayList<>();

    static {
        monitorPackages.add("xxx");
    }

    /**
     * 参数：
     * loader - 定义要转换的类加载器；如果是引导加载器，则为 null
     * className - 完全限定类内部形式的类名称和 The Java Virtual Machine Specification 中定义的接口名称。例如，"java/util/List"。
     * classBeingRedefined - 如果是被重定义或重转换触发，则为重定义或重转换的类；如果是类加载，则为 null
     * protectionDomain - 要定义或重定义的类的保护域
     * classfileBuffer - 类文件格式的输入字节缓冲区（不得修改）
     * 返回：
     * 一个格式良好的类文件缓冲区（转换的结果），如果未执行转换,则返回 null。
     * 抛出：
     * IllegalClassFormatException - 如果输入不表示一个格式良好的类文件
     */
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        boolean needModify = false;
        for (String p : monitorPackages) {
            if (className.startsWith(p)) {
                needModify = true;
                break;
            }
        }
        if (parentClassNames.contains(className)) {
            needModify = true;
        }
        if (!needModify) {
            return null;
        }
        try {
            // 借助JavaAssist工具，进行字节码插桩
            String classPackageName = className.replace("/", ".");
            ClassPool pool = ClassPool.getDefault();
            CtClass cc = pool.get(classPackageName);

            CtField[] declaredFields = cc.getDeclaredFields();
            CtClass superCc = cc.getSuperclass();
            // 处理子类，将当前子类属性存下来，当处理父类时，将父类的注解去除
            if (!parentClassNames.contains(className) && !superCc.getName().equals("java.lang.Object")) {
                List<String> subFieldNames = new ArrayList<>();
                for (int i = 0; i < declaredFields.length; i++) {
                    String name = declaredFields[i].getName();
                    subFieldNames.add(name);
                }
                parentClassNames.add(superCc.getName().replace(".", "/"));
                subFields.put(superCc.getName(), subFieldNames);
                return cc.toBytecode();
            }

            // 处理父类属性
            List<String> subFieldNames = subFields.get(classPackageName);
            if (null == subFieldNames) {
                return null;
            }
            for (int i = 0; i < declaredFields.length; i++) {
                CtField declaredField = declaredFields[i];
                if (subFieldNames.contains(declaredField.getName())) {
                    cc.removeField(declaredField);
                    CtField ctField = new CtField(declaredField.getType(), declaredField.getName(), cc);
                    cc.addField(ctField);
                }
            }
            return cc.toBytecode();
        } catch (Exception ex) {
            ex.printStackTrace();
            return null;
        }
    }
}
